import gym
from citywalk import PolicyIteration, print_agent, ValueIteration

env = gym.make('FrozenLake-v1')
env = env.unwrapped
env.render()

holes = set()
ends = set()
for s in env.P:
    for a in env.P[s]:
        for _s in env.P[s][a]:
            (_p, _ns, r, done) = _s
            if r == 1.0:
                ends.add(_ns)
            elif done:
                holes.add(_ns)
holes = holes - ends
print("冰洞：", holes)
print("目标：", ends)
for a in env.P[14]:
    print(a, env.P[14][a])

action_meaning = ['←', '↓', '→', '↑']
theta = 1e-5
gamma = 0.9
agent = PolicyIteration(env, theta, gamma)
agent.policy_iteration()
print_agent(agent, action_meaning, list(ends), list(holes))

agent = ValueIteration(env, theta, gamma)
agent.policy_iteration()
print_agent(agent, action_meaning, list(ends), list(holes))